﻿using Fst.SemiRing;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;

namespace Fst.Operations
{
    /// <summary>
    /// Extend an Fst to a single final state and undo operations.
    /// 
    /// @author John Salatas <jsalatas@users.sourceforge.net>
    /// </summary>
    public class ExtendFinal
    {
        /**
        /// Default Contructor
        */
        private ExtendFinal()
        {
        }

        /**
        /// Extends an Fst to a single final state.
        /// 
        /// It adds a new final state with a 0.0 (Semiring's 1) final wight and
        /// connects the current final states to it using epsilon transitions with
        /// weight equal to the original final state's weight.
        /// 
        /// @param fst the Fst to extend
         */
        public static void apply(Fst fst) 
        {
            Semiring semiring = fst.getSemiring();
            List<State> fStates = new List<State>();

            int numStates = fst.getNumStates();
            for (int i = 0; i < numStates; i++) {
                State s = fst.getState(i);
                if (s.getFinalWeight() != semiring.zero()) {
                    fStates.Add(s);
                }
            }

            // Add a new single final
            State newFinal = new State(semiring.one());
            fst.addState(newFinal);
            foreach (State s in fStates) 
            {
                // add epsilon transition from the old final to the new one
                s.addArc(new Arc(0, 0, s.getFinalWeight(), newFinal));
                // set old state's weight to zero
                s.setFinalWeight(semiring.zero());
            }
        }

        /**
        /// Undo of the extend operation
         */
        public static void undo(Fst fst)
        {
            State f = null;
            int numStates = fst.getNumStates();
            for (int i = 0; i < numStates; i++)
            {
                State s = fst.getState(i);
                if (s.getFinalWeight() != fst.getSemiring().zero())
                {
                    f = s;
                    break;
                }
            }

            if (f == null)
            {
                Trace.WriteLine("Final state not found.");
                return;
            }
            for (int i = 0; i < numStates; i++)
            {
                State s = fst.getState(i);
                for (int j = 0; j < s.getNumArcs(); j++)
                {
                    Arc a = s.getArc(j);
                    if (a.getIlabel() == 0 && a.getOlabel() == 0
                            && a.getNextState().getId() == f.getId())
                    {
                        s.setFinalWeight(a.getWeight());
                    }
                }
            }
            fst.deleteState(f);
        }

    }
}
